
import logging
from pprint import pformat
import sys 
import argparse

import torch

from lerobot.policies.factory import make_policy
from lerobot.utils.random_utils import set_seed
from lerobot.configs import parser
from lerobot.configs.train import TrainPipelineConfig

def get_teacher_model(model_path, dataset, bs):
    temp = sys.argv[1:]
    sys.argv[1:] = ['--config_path=' + model_path, '--resume=true', f'--batch_size={bs}']
    policy = compression(dataset)
    sys.argv[1:] = temp
    return policy

@parser.wrap()
def compression(cfg: TrainPipelineConfig, dataset):

    cfg.validate()
    logging.info(pformat(cfg.to_dict()))

    if cfg.seed is not None:
        set_seed(cfg.seed)

    # Check device is available
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = True

    logging.info("Creating policy")
    policy = make_policy(
        cfg=cfg.policy,
        ds_meta=dataset.meta,
    )
    policy.eval()
    return policy 

if __name__ == "__main__":
    # 创建参数解析器
    parser = argparse.ArgumentParser(
        description='处理模型配置文件（必填模型路径参数）',
        epilog='示例: python script.py --model_path=/path/to/train_config.json'
    )
    parser.add_argument('--model_path', type=str, help='教师模型训练配置文件路径')
    
    # 解析参数
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    args = parser.parse_args()
    policy = get_teacher_model(args.model_path)


